백준 온라인 저지 ▸ 10986

나머지 합
III

문제

수 N개 A₁, A₂, ..., A_N이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.

즉, A_i + ... + A_j (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 10⁶, 2 ≤ M ≤ 10³)

둘째 줄에 N개의 수 A₁, A₂, ..., A_N이 주어진다. (0 ≤ A_i ≤ 10⁹)

출력

첫째 줄에 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 출력한다.

예제 입력 1

5 3
1 2 3 1 2

예제 출력 1

풀이

코드

from math import comb
 
 
n, m = map(int, input().split())
a = list(map(int, input().split()))
 
prefix_sum = [e % m for e in a]
for i in range(1, n):
    prefix_sum[i] = (prefix_sum[i] + prefix_sum[i - 1]) % m
 
array = [0] * m
array[0] = 1
for p in prefix_sum:
    array[p] += 1
 
count = 0
for e in array:
    count += comb(e, 2)
 
print(count)